from glob import glob
import os
import json

import numpy as np
from torch.utils.data import Dataset
from torchvision import transforms

from utils.augmentation import Augmentation, cropCenter
from utils.basic_cache import BasicCache

def null_norm(x):
    return x

class Spacenet6Dataset(Dataset):
    def __init__(self, config) -> None:
        super()
        
        self.crop_size = config.crop if isinstance(config.crop, (int, float)) else None
        # 异源对的裁剪尺寸，默认crop_size_a大，作为搜索patch的尺寸，crop_size_b作为模板patch的尺寸
        self.crop_size_a = config.crop_a if isinstance(config.crop_a, (int, float)) else None
        self.crop_size_b = config.crop_b if isinstance(config.crop_b, (int, float)) else None
        
        self.named = config.named if isinstance(config.named, bool) else False
        self.stretch_contrast = config.stretch_contrast if isinstance(config.stretch_contrast, bool) else False
        self.return_all = config.return_all if isinstance(config.return_all, bool) else False
        # self.toDb = config.toDb if isinstance(config.toDb, bool) else False
        self.noise = config.noise if isinstance(config.noise, bool) else False
        self.zca = config.zca if isinstance(config.zca, bool) else False
        # 单源，仅输出一个可见光图像
        self.single_domain = config.single_domain if isinstance(config.single_domain, bool) else False
        # 没有该选项
        self.full_size = config.full_size if isinstance(config.full_size, bool) else False

        self.shift_range = config.shift_range if isinstance(config.shift_range, (list, tuple)) else [5, 15]

        # 读取角点文件
        self.ids2id_json = config.ids2id_json
        with open(self.ids2id_json) as fp:
            self.pt_set = json.load(fp)["points"]

        # 如果指定了缓存，那么我们会将补丁保存到本地磁盘的某个位置，以防止需要一直重新加载它们
        self.cache_dir = config.cache_dir if isinstance(config.cache_dir, str) else None
        self.cache_size = config.cache_size if isinstance(config.cache_size, (int, float)) else 0

        if self.cache_dir is not None:
            self.cache = BasicCache(self.cache_dir, size=self.cache_size, scheme="fill", clear=False, overwrite=False)
        else:
            self.cache = None

        if "sar" in config.normalize:
            self.sar_norm = transforms.Normalize(mean=[config.normalize.sar[0]], std=[config.normalize.sar[1]])
        else:
            self.sar_norm = null_norm

        if "opt" in config.normalize:
            self.opt_norm = transforms.Normalize(mean=[config.normalize.opt[0]], std=[config.normalize.opt[1]])
        else:
            self.opt_norm = null_norm
        
        print(f"")

    def __len__(self):
        return len(self.pt_set)

    # idx --> point_info
    def _index_to_point(self, idx: int) -> dict:
        pt_info = self.pt_set[idx]
        return pt_info

    def _load_and_label(self, index, drop=False):
        '''
        根据index返回经过归一化的sar和可见光图像对，
        返回为
        {"SAR" : sar块, 
         "OPT" : 可见光块, 
         "Y" : 0/1匹配标签, 
         "INFO" : 信息:
         }
        '''
        # data = self._try_cache(index)

        # 如果我们无法加载所需的补丁程序，至少要确保我们将回退缓存在原始密钥下，这样我们就不会再次循环
        # If we can't load the required patch, at least make sure we cache the fallback under the original key so we don't loop on this again
        # cache_key = self._index_to_point(index)[1].wkt
        
        if not data:
            data = {"SAR": None, "OPT": None, "Y": None, "INFO": None}

        # HACK：如果有太多无效补丁，会导致加载时间过长
        while (data["OPT"] is None) and (data["SAR"] is None):
            # 角点
            df_city, point = self._index_to_point(index)

            # 某一位置的异源图像窗口
            group = self.windows[df_city].groupby("wkt").get_group(point.wkt)

            imgs, idxs = self._get_group_patches(group, opt_transform=partial(self._normalize_scale, in_range=(0,255)), sar_transform=self._toDb)

            if len(idxs["OPT"]) > 0 and len(idxs["SAR"]) > 0:
                data["OPT"] = imgs["OPT"][0]
                data["SAR"] = imgs["SAR"][0]
                data["INFO"] = {"SAR": idxs["SAR"][0], "OPT": idxs["OPT"][0], "city": df_city, "supervised": group.super.values[0], "WKT": cache_key}
                # Urban Atlas is always corresponding on a patch level. Use sub patches for negatives
                data["Y"] = np.ones(1)

                # Try and cache the data point
                if self.cache is not None:
                    self.cache[cache_key] = data
            else:
                if drop:
                    break

        pass
    
class Spacenet6SiameseTriplet(Spacenet6Dataset):
    def multivariate_gaussian(self, pos, mu, Sigma):
        n = mu.shape[0]
        Sigma_det = np.linalg.det(Sigma)
        Sigma_inv = np.linalg.inv(Sigma)
        N = np.sqrt((2*np.pi)**n * Sigma_det)
        # This einsum call calculates (x-mu)T.Sigma-1.(x-mu) in a vectorized
        # way across all the input variables.
        fac = np.einsum('...k,kl,...l->...', pos-mu, Sigma_inv, pos-mu)
        return np.exp(-fac / 2) / N

    def get_gt_heatmap(self, shift_x=0, shift_y=0, w=64, h=64, sigma=1):
        '''
        返回w,h图中，偏移中心点shift_x,shift_y处为中心的sigma方差的高斯核的热力图
        '''
        x = int(w//2 + shift_x)
        y = int(h//2 + shift_y)

        if sigma is None:
            hm = np.zeros((h, w))
            hm[y, x] = 1
        else:
            X = np.linspace(0, w-1, w)
            Y = np.linspace(0, h-1, h)
            X, Y = np.meshgrid(X, Y)

            mu = np.array([x, y])
            Sigma = np.array([[sigma , 0], [0,  sigma]])
            pos = np.empty(X.shape + (2,))
            pos[:, :, 0] = X
            pos[:, :, 1] = Y
            hm = self.multivariate_gaussian(pos, mu, Sigma)

        return hm[np.newaxis, :, :]
    
    def __getitem__(self, index):
        # 图像增广
        if self.augmentor:
            self.augmentor.refresh_random_state()

        # 获取源图像块和标签信息
        img_sar, img_opt, y, img_info = self._load_and_label(index)

        # 设置图像增强
        pass

        # 单源
        pass

        assert self.crop_size_a <= img_sar.shape[1], "The input image is too small to crop"
        assert self.crop_size_b <= img_opt.shape[1], "The input image is too small to crop"

        max_shift = min(min(self.crop_size_a) // 4, min(self.crop_size_b) // 4) # 最大位移一般是32
        shift_x = (2 * np.random.randint(2) - 1) * (np.random.randint(max_shift) + 1)
        shift_y = (2 * np.random.randint(2) - 1) * (np.random.randint(max_shift) + 1)

        if self.crop_size_a > self.crop_size_b:
            max_shift_x = min(max_shift - np.abs(shift_x//2), max_shift)
            max_shift_y = min(max_shift - np.abs(shift_y//2), max_shift)
            shift_x_s = (2*np.random.randint(2) - 1)*(np.random.randint(max_shift_x))
            shift_y_s = (2*np.random.randint(2) - 1)*(np.random.randint(max_shift_y))

            search_img = np.ascontiguousarray(cropCenter(img_sar, (self.crop_size_a, self.crop_size_a), (shift_x_s, shift_y_s)))
            template_img = np.ascontiguousarray(cropCenter(img_opt, (self.crop_size_a, self.crop_size_a), (shift_x_s, shift_y_s)))
            search_hard = np.ascontiguousarray(cropCenter(img_sar, (self.crop_size_b, self.crop_size_b), (shift_x, shift_y)))
            template_hard = np.ascontiguousarray(cropCenter(img_opt, (self.crop_size_b, self.crop_size_b), (shift_x, shift_y)))

            # 归一化
            pass

        else:
            max_shift_x = min(max_shift - np.abs(shift_x//2), max_shift)
            max_shift_y = min(max_shift - np.abs(shift_y//2), max_shift)
            shift_x_s = (2*np.random.randint(2) - 1)*(np.random.randint(max_shift_x))
            shift_y_s = (2*np.random.randint(2) - 1)*(np.random.randint(max_shift_y))

            search_img = np.ascontiguousarray(cropCenter(img_opt, (self.crop_size_b, self.crop_size_b), (shift_x_s, shift_y_s)))
            template_img = np.ascontiguousarray(cropCenter(img_sar, (self.crop_size_b, self.crop_size_b), (shift_x_s, shift_y_s)))
            search_hard = np.ascontiguousarray(cropCenter(img_opt, (self.crop_size_a, self.crop_size_a), (shift_x, shift_y)))
            template_hard = np.ascontiguousarray(cropCenter(img_sar, (self.crop_size_a, self.crop_size_a), (shift_x, shift_y)))

            # 归一化
            pass
